package com.example.demo.camera2;

import android.graphics.Bitmap;
import android.graphics.Color;
import android.graphics.ImageFormat;
import android.hardware.camera2.CameraCharacteristics;
import android.hardware.camera2.DngCreator;
import android.hardware.camera2.TotalCaptureResult;
import android.media.Image;
import android.os.Environment;
import android.os.SystemClock;
import android.util.Log;

import java.io.File;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.OutputStream;
import java.io.OutputStreamWriter;
import java.nio.ByteBuffer;
import java.nio.FloatBuffer;
import java.nio.ShortBuffer;
import java.nio.channels.Channels;
import java.nio.channels.WritableByteChannel;
import java.text.SimpleDateFormat;
import java.util.Date;

public class ImageUtil {
    private static final String TAG = "ImageUtil";

    public String saveImage(Image img) throws IOException {
        long timestamp = img.getTimestamp();
        File output = getOutputImageFile(img.getFormat(), timestamp);
        try (FileOutputStream out = new FileOutputStream(output)) {
            switch(img.getFormat()) {
                case ImageFormat.JPEG: {
                    writeJpegImage(img, out);
                    break;
                }
                case ImageFormat.YUV_420_888: {
                    writeYuvImage(img, out);
                    break;
                }
                case ImageFormat.RAW_SENSOR: {
                    writeDngImage(img, out);
                    break;
                }
                case ImageFormat.RAW10: {
                    TLog.e("RAW10 saving not implemented");
                    break;
                }
                case ImageFormat.DEPTH16: {
                    writeDepth16Image(img, out);
                    break;
                }
                case ImageFormat.DEPTH_POINT_CLOUD: {
                    writeDepthPointImage(img, out);
                    break;
                }
            }
        }
        return output.getName();
    }

    private void writeDngImage(Image img, OutputStream out) throws IOException {
//        if (img.getFormat() != ImageFormat.RAW_SENSOR) {
//            throw new IOException(
//                    String.format("Unexpected Image format: %d, expected ImageFormat.RAW_SENSOR",
//                            img.getFormat()));
//        }
//        long timestamp = img.getTimestamp();
//        if (mCurrentCamera == null) {
//            TLog.e("No camera availble for camera info, not saving DNG (timestamp %d)",
//                    timestamp);
//            throw new IOException("No camera info available");
//        }
//        TotalCaptureResult result = mCurrentCamera.getResultAt(timestamp);
//        if (result == null) {
//            TLog.e("No result matching raw image found, not saving DNG (timestamp %d)",
//                    timestamp);
//            throw new IOException("No matching result found");
//        }
//        CameraCharacteristics info = mCurrentCamera.getCharacteristics();
//        try (DngCreator writer = new DngCreator(info, result)) {
//            writer.writeImage(out, img);
//        }
    }

    private void writeJpegImage(Image img, OutputStream out) throws IOException {
        if (img.getFormat() != ImageFormat.JPEG) {
            throw new IOException(
                    String.format("Unexpected Image format: %d, expected ImageFormat.JPEG",
                            img.getFormat()));
        }
        WritableByteChannel outChannel = Channels.newChannel(out);
        ByteBuffer jpegData = img.getPlanes()[0].getBuffer();
        jpegData.rewind();
        outChannel.write(jpegData);
    }

    private void writeYuvImage(Image img, OutputStream out)
            throws IOException {
        if (img.getFormat() != ImageFormat.YUV_420_888) {
            throw new IOException(
                    String.format("Unexpected Image format: %d, expected ImageFormat.YUV_420_888",
                            img.getFormat()));
        }
        WritableByteChannel outChannel = Channels.newChannel(out);
        for (int plane = 0; plane < 3; plane++) {
            Image.Plane colorPlane = img.getPlanes()[plane];
            ByteBuffer colorData = colorPlane.getBuffer();
            int subsampleFactor = (plane == 0) ? 1 : 2;
            int colorW = img.getWidth() / subsampleFactor;
            int colorH = img.getHeight() / subsampleFactor;
            colorData.rewind();
            colorData.limit(colorData.capacity());
            if (colorPlane.getPixelStride() == 1) {
                // Can write contiguous rows
                for (int y = 0, rowStart = 0; y < colorH;
                        y++, rowStart += colorPlane.getRowStride()) {
                    colorData.limit(rowStart + colorW);
                    colorData.position(rowStart);
                    outChannel.write(colorData);
                }
            } else {
                // Need to pack rows
                byte[] row = new byte[(colorW - 1) * colorPlane.getPixelStride() + 1];
                byte[] packedRow = new byte[colorW];
                ByteBuffer packedRowBuffer = ByteBuffer.wrap(packedRow);
                for (int y = 0, rowStart = 0; y < colorH;
                        y++, rowStart += colorPlane.getRowStride()) {
                    colorData.position(rowStart);
                    colorData.get(row);
                    for (int x = 0, i = 0; x < colorW;
                            x++, i += colorPlane.getPixelStride()) {
                        packedRow[x] = row[i];
                    }
                    packedRowBuffer.rewind();
                    outChannel.write(packedRowBuffer);
                }
            }
        }
    }

    /**
     * Save a 16-bpp depth image as a false-color PNG
     */
    private void writeDepth16Image(Image img, OutputStream out) throws IOException {
        if (img.getFormat() != ImageFormat.DEPTH16) {
            throw new IOException(
                    String.format("Unexpected Image format: %d, expected ImageFormat.DEPTH16",
                            img.getFormat()));
        }
        int w = img.getWidth();
        int h = img.getHeight();
        int rowStride = img.getPlanes()[0].getRowStride() / 2; // in shorts
        ShortBuffer y16Data = img.getPlanes()[0].getBuffer().asShortBuffer();

        Bitmap rgbImage = convertDepthToFalseColor(y16Data, w, h, rowStride, /*scale*/ 1);
        rgbImage.compress(Bitmap.CompressFormat.PNG, 100, out);
        rgbImage.recycle();
    }

    // This saves a text file of float values for a point cloud
    private void writeDepthPointImage(Image img, OutputStream out) throws IOException {
        if (img.getFormat() != ImageFormat.DEPTH_POINT_CLOUD) {
            throw new IOException(
                    String.format("Unexpected Image format: %d, expected ImageFormat.DEPTH16",
                            img.getFormat()));
        }
        FloatBuffer pointList = img.getPlanes()[0].getBuffer().asFloatBuffer();
        int pointCount = pointList.limit() / 3;
        OutputStreamWriter writer = new OutputStreamWriter(out);
        for (int i = 0; i < pointCount; i++) {
            String pt = String.format("%f, %f, %f\n",
                    pointList.get(), pointList.get(),pointList.get());
            writer.write(pt, 0, pt.length());
        }
    }

    File getOutputImageFile(int type, long timestamp){
        // To be safe, you should check that the SDCard is mounted
        // using Environment.getExternalStorageState() before doing this.

        String state = Environment.getExternalStorageState();
        if (!Environment.MEDIA_MOUNTED.equals(state)) {
            return null;
        }

        File mediaStorageDir = new File(Environment.getExternalStoragePublicDirectory(
                  Environment.DIRECTORY_DCIM), "TestingCamera2");
        // This location works best if you want the created images to be shared
        // between applications and persist after your app has been uninstalled.

        // Create the storage directory if it does not exist
        if (!mediaStorageDir.exists()){
            if (!mediaStorageDir.mkdirs()){
                TLog.e("Failed to create directory for pictures/video");
                return null;
            }
        }

        // Create a media file name

        // Find out time now in the Date and boottime time bases.
        long nowMs = new Date().getTime();
        long nowBootTimeNs = SystemClock.elapsedRealtimeNanos();

        // Convert timestamp from boottime time base to the Date timebase
        // Slightly approximate, but close enough
        final long NS_PER_MS = 1000000l;
        long timestampMs = (nowMs * NS_PER_MS - nowBootTimeNs + timestamp) / NS_PER_MS;

        String timeStamp = new SimpleDateFormat("yyyy-MM-dd_HH-mm-ss_SSS").
                format(new Date(timestampMs));
        File mediaFile = null;
        switch(type) {
            case ImageFormat.JPEG:
                mediaFile = new File(mediaStorageDir.getPath() + File.separator +
                        "IMG_"+ timeStamp + ".jpg");
                break;
            case ImageFormat.YUV_420_888:
                mediaFile = new File(mediaStorageDir.getPath() + File.separator +
                        "IMG_"+ timeStamp + ".yuv");
                break;
            case ImageFormat.RAW_SENSOR:
                mediaFile = new File(mediaStorageDir.getPath() + File.separator +
                        "IMG_"+ timeStamp + ".dng");
                break;
            case ImageFormat.RAW10:
                mediaFile = new File(mediaStorageDir.getPath() + File.separator +
                        "IMG_"+ timeStamp + ".raw10");
                break;
            case ImageFormat.DEPTH16:
                mediaFile = new File(mediaStorageDir.getPath() + File.separator +
                        "IMG_"+ timeStamp + "_depth.png");
                break;
            case ImageFormat.DEPTH_POINT_CLOUD:
                mediaFile = new File(mediaStorageDir.getPath() + File.separator +
                        "IMG_"+ timeStamp + "_depth_points.txt");
                break;
            default:
                mediaFile = new File(mediaStorageDir.getPath() + File.separator +
                        "IMG_"+ timeStamp + ".unknown");
                TLog.e("Unknown image format for saving, using .unknown extension: " + type);
                break;
        }

        return mediaFile;
    }

    /**
     * Convert depth16 buffer into a false-color RGBA Bitmap, scaling down
     * by factor of scale
     */
    private Bitmap convertDepthToFalseColor(ShortBuffer depthBuffer, int w, int h,
            int stride, int scale) {
        short[] yRow = new short[w];
        int[] imgArray = new int[w * h];
        w = w / scale;
        h = h / scale;
        stride = stride * scale;
        for (int y = 0, j = 0, rowStart = 0; y < h; y++, rowStart += stride) {
            // Align to start of nearest-neighbor row
            depthBuffer.position(rowStart);
            depthBuffer.get(yRow);
            for (int x = 0, i = 0; x < w; x++, i += scale, j++) {
                short y16 = yRow[i];
                int r = y16 & 0x00FF;
                int g = (y16 >> 8) & 0x00FF;
                imgArray[j] = Color.rgb(r, g, 0);
            }
        }
        return Bitmap.createBitmap(imgArray, w, h, Bitmap.Config.ARGB_8888);
    }
}

